from __future__ import annotations

import gymnasium as gym
import numpy as np

from driving_gym.data.data_provider import Snapshot
from driving_gym.environment.agent.obs.base_obs import BaseObs
from driving_gym.environment.scenario.base_scenario import BaseScenario
from driving_gym.misc.util import override
from driving_gym.simulation.adapter_interface import AdapterInterface


class SensorSemanticObs(BaseObs):
    def __init__(self, config: dict, adapter: AdapterInterface):
        super().__init__(config, adapter)

        self.gnss_source = config.get("gnss_source", "gnss")
        self.imu_source = config.get("imu_source", "imu")
        self.speed_source = config.get("speed_source", "speedometer")

    @override(BaseObs)
    def get_obs(self, snapshot: Snapshot, scenario: BaseScenario) -> np.ndarray:
        """Process the semantic observation from the snapshot data.

        Args:
            snapshot (Snapshot): One frame of the simulation data

        Returns:
            np.ndarray: [latitude, longitude, altitude, *accelerometer, *gyroscope, compass, speed]
        """
        gnss_data = snapshot.data[self.actor_id]["sensors"][self.gnss_source]
        imu_data = snapshot.data[self.actor_id]["sensors"][self.imu_source]
        speed_data = snapshot.data[self.actor_id]["sensors"][self.speed_source]

        # NaN could exist due to sensor bugs (Typically compass in IMU)
        return np.nan_to_num(np.concatenate([gnss_data, imu_data, speed_data]))

    @override(BaseObs)
    def get_observation_space(self) -> gym.spaces.Box:
        return gym.spaces.Box(low=-np.inf, high=np.inf, shape=(11,), dtype=np.float64)
